import "DPI-C" function void psram_read(input int addr, output int data);
import "DPI-C" function void psram_write(input int addr, input int data, input int size);

module psram(
  input sck,
  input ce_n,
  inout [3:0] dio
);

  wire reset = ce_n;
  wire [3:0] din, dout;

  typedef enum [2:0] { cmd_t, cmd4_t, addr_t, rwait_t, rdata_t, wdata_t, err_t } state_t;
  reg [2:0] curr_state;
  reg [2:0] next_state;

  reg [3:0]  counter;
  reg [7:0]  cmd;
  reg [23:0] addr;
  reg [31:0] wdata;
  reg [31:0] rdata;
  reg        wen;
  reg [2:0]  wsize;

  reg QPI = 1'b0;

  always @(posedge sck or posedge reset) begin
    if (reset && !QPI) curr_state <= cmd_t;
    else if (reset && QPI) curr_state <= cmd4_t;
    else curr_state <= next_state;
  end
  always @(*) begin
    case (curr_state)
      cmd_t  : next_state = (counter == 4'd7)                    ? cmd4_t  : cmd_t;
      cmd4_t : next_state = (counter == 4'd1)                    ? addr_t  : cmd4_t;
      addr_t : next_state = (cmd     != 8'hEB && cmd   != 8'h38) ? err_t   :
                            (counter != 4'd5)                    ? addr_t  :
                            (cmd     == 8'hEB)                   ? rwait_t : wdata_t;
      rwait_t: next_state = (counter == 4'd6)                    ? rdata_t : rwait_t;
      rdata_t: next_state = rdata_t;
      wdata_t: next_state = wdata_t;
      default: begin
        next_state = curr_state;
          $fwrite(32'h80000002, "Assertion failed: Unsupported command `%xh`, only support `EBh` read and `38h` write command\n", cmd);
          $fatal;
      end
    endcase
  end

  always @(posedge sck or posedge reset) begin
    if (reset) counter <= 4'd0;
    else begin
      case (curr_state)
        cmd_t  :   counter <= (counter < 4'd7) ? counter + 4'd1 : 4'd0;
        cmd4_t :   counter <= (counter < 4'd1) ? counter + 4'd1 : 4'd0;
        addr_t :   counter <= (counter < 4'd5) ? counter + 4'd1 : 4'd0;
        rwait_t:   counter <= (counter < 4'd6) ? counter + 4'd1 : 4'd0;
        wdata_t:   counter <= (counter < 4'd7) ? counter + 4'd1 : 4'd8;
        default:   counter <= counter + 4'd1;
      endcase
    end
  end

  always @(negedge sck) begin
    if (cmd == 8'hCC) QPI = 1'b1;
  end

  // get command
  always @(posedge sck or posedge reset) begin
    if (reset)                     cmd <= 8'd0;
    else if (curr_state == cmd_t)  cmd <= {cmd[6:0], din[0]};
    else if (curr_state == cmd4_t) cmd <= {cmd[3:0], din};
  end

  // get address
  always @(posedge sck or posedge reset) begin
    if (reset) addr <= 24'd0;
    else if (curr_state == addr_t)
      addr <= { addr[19:0], din };
  end

  // get writing data
  always @(posedge sck or posedge reset) begin
    if (reset) wdata <= 32'd0;
    else if (curr_state == wdata_t && counter < 4'd2)                   wdata <= { wdata[27:0],  din };
    else if (curr_state == wdata_t && counter < 4'd4 && counter >=4'd2) wdata <= { wdata[27:8],  din, wdata[7:0] };
    else if (curr_state == wdata_t && counter < 4'd6 && counter >=4'd4) wdata <= { wdata[27:16], din, wdata[15:0] };
    else if (curr_state == wdata_t && counter < 4'd8 && counter >=4'd6) wdata <= { wdata[27:24], din, wdata[23:0] };
  end
  // write enable
  always @(negedge sck) begin
    if (curr_state == wdata_t && counter == 4'd1)      begin wen <= 1'b1; wsize <= 3'd1; end
    else if (curr_state == wdata_t && counter == 4'd3) begin wen <= 1'b1; wsize <= 3'd2; end
    else if (curr_state == wdata_t && counter == 4'd7) begin wen <= 1'b1; wsize <= 3'd4; end
    else begin wen <= 1'b0;  wsize <= 3'd0; end
  end
  // write data
  always @(negedge sck) begin
    if (wen) psram_write({ 8'd0, addr }, wdata, { 29'd0, wsize });
  end

  // read data
  always @(posedge sck) begin
    if (curr_state == rwait_t && counter == 0) psram_read({ 8'd0, addr }, rdata);
  end

  assign dio = dout;
  assign dout = (curr_state != rdata_t) ? 4'bz :
                (counter == 0) ? rdata[7:4]    :
                (counter == 1) ? rdata[3:0]    :
                (counter == 2) ? rdata[15:12]   :
                (counter == 3) ? rdata[11:8]  :
                (counter == 4) ? rdata[23:20]  :
                (counter == 5) ? rdata[19:16]  :
                (counter == 6) ? rdata[31:28]  :
                (counter == 7) ? rdata[27:24]  : 4'bz;
  assign din = dio;

endmodule
